class Index:
    def __init__(self):
        self.map = dict()

    def addKeyWord(self, keyword: str, filename: str, hit: int):
        """
        添加元素
        :param keyword: 关键词
        :param filename: 文件名
        :param hit: 位置
        :return: null
        """
        if filename not in self.map:
            self.map[filename] = {}
        if keyword not in self.map[filename]:
            self.map[filename][keyword] = []
        self.map[filename][keyword].append(hit)


class Vector:
    @staticmethod
    def vector(keyword: str, filename: str, index: Index, __filter=None) -> (list, list):
        if __filter is None:
            __filter = ['the', 'a', 'an', 'in', 'on', 'to']
        file_vector, search_vector = [], []
        for word in index.map[filename]:
            if word in __filter:
                continue
            file_vector.append(len(index.map[filename][keyword]))
            search_vector.append(1 if keyword == word else 0)
        return file_vector, search_vector


class InvertedIndex:
    def __init__(self):
        self.map = dict()

    def addKeyWord(self, keyword: str, filename: str, hit: int):
        if keyword not in self.map:
            self.map[keyword] = {}
        if filename not in self.map[keyword]:
            self.map[keyword][filename] = []
        self.map[keyword][filename].append(hit)

    def toInverted(self, index: Index):
        for filename in index.map:
            for keyword in index.map[filename]:
                for hit in index.map[filename][keyword]:
                    self.addKeyWord(keyword, filename, hit)
        return self

    def search(self, keyword: str) -> dict:
        """
        寻找关键词
        :param keyword: 关键词
        :return: 返回字典{filename=>list(hit)}
        """
        if keyword in self.map:
            return self.map[keyword]
        return dict()


class OutColors:
    CRED = '\033[91m'
    CEND = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


if __name__ == "__main__":
    """
    test
    """
    a = InvertedIndex()
    a.addKeyWord("1", "a", 1)
    a.addKeyWord("1", "a", 2)
    a.addKeyWord("1", "a", 3)
    a.addKeyWord("1", "a", 4)
    a.addKeyWord("2", "a", 4)
    a.addKeyWord("2", "a", 3)
    a.addKeyWord("2", "a", 2)
    a.addKeyWord("2", "a", 1)
    print(a.map)
    print("=" * 5, "InvertedIndex", "=" * 5)
    b = Index()
    c = InvertedIndex()
    b.addKeyWord("You", "1.txt", 1)
    b.addKeyWord("are", "1.txt", 2)
    b.addKeyWord("a", "1.txt", 3)
    b.addKeyWord("pig", "1.txt", 4)
    b.addKeyWord("mjj", "2.txt", 1)
    b.addKeyWord("is", "2.txt", 2)
    b.addKeyWord("a", "2.txt", 3)
    b.addKeyWord("man", "2.txt", 4)
    b.addKeyWord("a", "2.txt", 5)
    c.toInverted(b)
    print(b.map)
    print(c.map)
    print("=" * 5, "index", "=" * 5)
